package core

import (
	"huwhy.cn/go-spider/model"
	"huwhy.cn/go-spider/parser"
)

type Spider struct {
	reqWorkers  []RequestWorker
	respWorkers []ResponseWorker
	taskQueue   *model.SyncQueue
	resultQueue *model.SyncQueue
	storage     model.Storage
	parsers     map[string]parser.Parser
	quit        chan struct{}
}

func (s *Spider) AddTask(req *model.Request) {
	s.taskQueue.Push(req)
}

func (s *Spider) AddResult(resp *model.Response) {
	s.resultQueue.Push(resp)
}

func (s *Spider) SetStorage(storage model.Storage) {
	s.storage = storage
}

func (s *Spider) RegParser(_parser parser.Parser) {
	s.parsers[_parser.Name()] = _parser
}

func (s *Spider) getParser(name string) parser.Parser {
	if p, ok := s.parsers[name]; ok {
		return p
	}
	return nil
}

func (s *Spider) Start() {
	for i := 0; i < len(s.reqWorkers); i++ {
		go s.reqWorkers[i].Run()
	}
	for i := 0; i < len(s.respWorkers); i++ {
		go s.respWorkers[i].Run()
	}
	<-s.quit
	for i := 0; i < len(s.reqWorkers); i++ {
		s.reqWorkers[i].Stop()
	}
	for i := 0; i < len(s.respWorkers); i++ {
		s.respWorkers[i].Stop()
	}
}

func (s *Spider) Init(workerNum int) {
	s.taskQueue = model.NewSyncQueue()
	s.resultQueue = model.NewSyncQueue()
	s.storage = &model.EmptyStorage{}
	s.parsers = make(map[string]parser.Parser)
	for i := 0; i < workerNum; i++ {
		s.reqWorkers = append(s.reqWorkers, RequestWorker{
			id:     i + 1,
			state:  true,
			spider: s,
		})
	}
	for i := 0; i < workerNum; i++ {
		s.respWorkers = append(s.respWorkers, ResponseWorker{
			id:     i + 1,
			state:  true,
			spider: s,
		})
	}
}
